ENH: add new function `one_hot` by NeilGirdhar · Pull Request #306 · data-apis/array-api-extra

NeilGirdhar · 2025-05-29T14:24:33Z

Fixes #305

Questions:

What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

lucascolley · 2025-05-30T10:23:41Z

RE dtype, I think data-apis/array-api#848 will give us something slightly cleaner down the line.

In SciPy we have been using https://github.com/scipy/scipy/blob/main/scipy/_lib/_array_api.py#L399 with force_floating=True. Taking xp.empty(()).dtype is probably fine for now though.

NeilGirdhar · 2025-05-30T13:16:35Z

Yeah, I'm not 100% sure how those links will help in this case? We're not casting one thing to another kind, or promoting to float. We just want the default float dtype irrespective of what was passed in. You may want to consider giving names to:

xp.asarray(1j).dtype  # Default complex
xp.asarray(1).dtype  # Default int
xp.empty(()).dtype  # Default float

I'm not sure though, and these are one-liners.

lucascolley · 2025-05-30T13:29:17Z

I'm not 100% sure how those links will help in this case?

the point is that instead of writing

dtype = xp.empty(()).dtype
x = xp.zeros(..., dtype=dtype)

we would have

x = xp.zeros(...)
x = xp.astype(x, 'real floating')

NeilGirdhar · 2025-05-30T13:30:27Z

Ah, right. Okay. I guess you mean askind, right?

Also, could you help me solve the Dask errors? This is all foreign to me.

And how do I make an array with a non-concrete size? (x.size=None)

lucascolley · 2025-05-30T13:46:50Z

I guess you mean askind, right?

Nope, data-apis/array-api#848 suggests overloading the dtype parameter of astype for this. Feel free to comment over there if you disagree!

NeilGirdhar · 2025-05-30T14:46:33Z

suggests overloading the dtype parameter of astype for this. Feel free to comment over there if you disagree!

Oh, no worries, I don't have a strong opinion. Just trying to keep up with all the planned changes 😄

Did you see my edits? I could use some guidance with the Dask errors.

ghost · 2025-05-30T16:50:45Z

What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

to me this doesn't make much sense. Why shouldn't it be bool?

ghost · 2025-05-30T22:02:52Z

What should the default output dtype be? Currently, it's the default float type on Jax/PyTorch, or else float64. Should I try to get the default float type on the platform: xp.empty(()).dtype?

to me this doesn't make much sense. Why shouldn't it be bool?

I'd rather follow what the libraries are doing than force double conversion for delegated code. If you do that, most people would end up having to write their own one_hot method in order to avoid it.

In general, the reason it's not bool is because these values often serve as the inputs to machine learning algorithms.

ghost · 2025-05-31T02:55:11Z

@lucascolley Ready for your review

ghost · 2025-06-01T21:28:53Z

[EDIT] just realised that I have had an unsent review hanging for the last 3 days. Apologies.

What about x[..., None] == xp.arange(num_classes, dtype=x.dtype, device=_compact.device(x))?

that should be all you need to do to implement one_hot in a way that is both performant and Array API compliant.
(the numpy special case may be faster; needs benchmarking).

ghost · 2025-06-02T10:20:27Z

+    Parameters
+    ----------
+    x : array
+        An array with integral dtype and concrete size (``x.size`` cannot be `None`).


I think with the new algorithm none sizes are supported (needs a test though)

ghost · 2025-06-02T10:23:54Z

+        ]
+    # Delegate where possible.
+    if is_jax_namespace(xp):
+        assert is_jax_array(x)


Suggested change

assert is_jax_array(x)

We typically don't do this. It's needlessly expensive.

I agree that we should be consistent with the rest of the package, so I'll remove it. However, you may want to revisit this policy since with jitted code such calls are basically free. This is true for Jax and Pytorch, which are the backends that I added assertions with. Also, assertions disappear when optimizations are turned on.

The assertions have the benefit of improve type checking.

ghost

This is almost ready to go; a few nits below.

It's worth noting that the current implementation makes it impossible to write

symbols, idx = xp.unique_inverse(x)
xpx.one_hot(idx, symbols.size)

as a pattern to build a one-hot map of arbitrary symbols on Dask, unless you know in advance the maximum number of unique symbols.
This could be fixed in a follow-up but I'm unsure about real-life interest in it.

ghost · 2025-06-02T16:11:47Z

looks like there was a rebase hiccup

ghost

thanks both, looks great!

ghost · 2025-06-03T02:05:02Z

Thanks @lucascolley and @crusaderky for the thorough and quick reviews!

ghost

Only one nit; see above

ghost · 2025-06-03T09:05:49Z

looks like there was a rebase hiccup

ghost · 2025-06-03T09:46:29Z

thanks again!

lucascolley changed the title ~~Add one_hot~~ ENH: add new function one_hot May 29, 2025

lucascolley added enhancement New feature or request new function labels May 29, 2025

lucascolley added this to the 0.8.0 milestone May 29, 2025

lucascolley reviewed May 30, 2025

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

lucascolley reviewed May 30, 2025

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

Comment thread src/array_api_extra/_delegation.py

Comment thread src/array_api_extra/_delegation.py Outdated

lucascolley reviewed May 30, 2025

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

lucascolley reviewed May 31, 2025

View reviewed changes

Comment thread src/array_api_extra/_delegation.py

crusaderky reviewed Jun 1, 2025

View reviewed changes